#!/usr/bin/python
# VERSION 1.1 Peter Daldrop 28/1/2008
#changelog: 1.0 first version, 1.01 flexible lattice constants, 1.02 correct atom numbers for spheres
#developing: accepts .sph and modified (sypdbxtal.py).pdb files to circumvent need for .sph first but output is still sph
#1.1major changes to allow more options from command line, generated centered cubic lattices and 
import os,sys,string,math,random
from Numeric import *
from numpy import *
#print os.getcwd()
#This script calculates a box around the ligand molecule. To make the box identical in size and shape irrespective of the 
#ligand's orientation in space the torque of inertia is calculated and used for a coordinate transformation. The script
#needs two aruments, the first is the .sph input file containing the ligand atoms and the second is the output file containing the sphere lattice
####Defining functions
def sphconv(file_in):

#output file
    file_out = file_in+".pdb"
    f_in = open(file_in,'r')
    f_out = open(file_out , 'w')

    start = 'false'

    atom_type_dict ={7:'C', 3:'O', 4:'N', 6:'I'}
    

    for i in f_in.xreadlines():
        if start == 'true':
            #print i
            number, x, y, z, junk1, junk2, junk3, colour = string.split(string.strip(i))
            out_line = 'ATOM' + string.rjust(number, 7) + string.rjust(atom_type_dict[int(colour)], 3) + string.rjust('SPH', 6)+ string.rjust(number,6) + string.rjust(x[0:5],12) + string.rjust(y[0:6],8) + string.rjust(z[0:6],8) 
            f_out.write(out_line + '\n')
            f_out.write('TER\n') 
      
        elif i[0:7] == 'cluster':
            start = 'true'
     
    f_out.close()
    f_in.close() 
   
   
  
def coordout(coord):
	out="  "+str(coord)+"00000"	
	out=string.split(out,".") #split at .
	out=out[0][-3:]+"."+out[1][:5] #refuse with right format
	return out #return correct string
###Start of script

def readcoords(file):
    n=1
    atoms=0
    xcoord=[0,] #create list for x coordinates
    ycoord=[0,] #create list for y coordinates
    zcoord=[0,] #create list for z coordinates
    atomcoord=[] #Variable to store coordinate array
    if dattype=="sph":
		print "Datafile is of .sph type"
		for i in file.readlines(): #for all lines
			if i[0]==" ": #true if line contains sphere coordinates
			   	  xcoord.append(float(i[6:14]))  #add x coordinate to list
			   	  ycoord.append(float(i[16:24])) #add y coordinate to list
				  zcoord.append(float(i[26:34])) #add z coordinate to list
				  atomcoord.append(array((xcoord[n], ycoord[n], zcoord[n]))) 
			      #creates list with all all atom's coordinates stored as array
				  n=n+1
			else:#if not sphere coordinates containing just write to output file
				out_file.write(i)
		return atomcoord
		print "Atom coordinates of ligand have been read in"
    elif dattype=="pdb":
		print "Datafile is of .pdb type"
		out_file.write("DOCK 5.2 ligand_atoms\npositive                       (1)\nnegative                       (2)\nacceptor                       (3)\ndonor                          (4)\nester_o                        (5)\namide_o                        (6)\nneutral                        (7)\ncluster\n")
		for i in data_file.readlines(): #for all lines
			if (i.startswith("ATOM") or i.startswith("HETATM")):
				atoms+=1			
				xcoord.append(float(i[31:38])) #add x coordinate to list
				ycoord.append(float(i[39:46])) #add y coordinate to list
				zcoord.append(float(i[47:54])) #add z coordinate to list
				atomcoord.append(array((xcoord[n], ycoord[n], zcoord[n])))
			
				n=n+1	
		return atomcoord
		print str(atoms)+" atoms read in from pdb file"


    else:
		print "Error! Datatype not recognized! Script aborted."
		sys.exit()

def svdtrans(coordlist):
    print "svd-transformation started...."
    coordlist=array(coordlist)#reformat coordlist
    molcent=array([mean(coordlist[:,0]),mean(coordlist[:,1]),mean(coordlist[:,2])])#calculate center
    centcord=coordlist-molcent #center coordinate list
    u,w,vt=linalg.svd(centcord)# perform svd transformation
    v=transpose(vt)  #v and vt will be used as retransformation matrix
    transc=dot(centcord,v)#transform coordinates
    print "done"
    return [molcent,v,transc] #return center of mass, transf matrix and transformed coords
    
def boxcalc(coords,rim):
    print "determining sphere box"
    #rim=1.5
    boxcorn=array([[max(coords[:,0])+rim,min(coords[:,0])-rim],[max(coords[:,1])+rim,min(coords[:,1])-rim],[max(coords[:,2])+rim,min(coords[:,2])-rim]]) #gets maximum atom distance from origin in X direction 
    return array(boxcorn)
def latticegen(box,netwidth,latticetype):
    #latticetype=2
    lattice=[]
    print "starting lattice generation"
    
    spans=[box[0,0]-box[0,1],box[1,0]-box[1,1],box[2,0]-box[2,1]]
    print "spans:"+str(spans)
    
    dims=[divmod(spans[0],netwidth)[0],divmod(spans[1],netwidth)[0],divmod(spans[2],netwidth)[0]]
    nets=[spans[0]/dims[0],spans[1]/dims[1],spans[2]/dims[2]]
    print type(latticetype)
    print nets
    if latticetype==1:
        print "using simple cubic lattice"
        tempx=box[0,1]
        while tempx <= box[0,0]+0.01:
            
            tempy=box[1,1]
            while tempy <=box[1,0]+0.01:
                tempz=box[2,1]
                while tempz<=box[2,0]+0.01:
                    lattice.append(array([tempx,tempy,tempz]))
                    tempz=tempz+nets[2]
                tempy=tempy + nets[1]
            tempx =tempx + nets[0]
        #print array(lattice)
        return lattice
    if latticetype==2:
        print "using centered cubic lattice"
        
        tempx=box[0,1]
        nx=0
        while tempx <= box[0,0]+0.01:
            ny=0
            tempy=box[1,1]
            while tempy <=box[1,0]+0.01:
                tempz=box[2,1]
                nz=0
                while tempz<=box[2,0]+0.01:
                    lattice.append(array([tempx,tempy,tempz]))
                    if ((nx-dims[0])*(ny-dims[1])*(nz-dims[2]))!=0:
                        lattice.append(array([tempx+0.5*nets[0],tempy+0.5*nets[1],tempz+0.5*nets[2]]))
                    tempz=tempz+nets[2]
                    nz+=1
                tempy=tempy + nets[1]
                ny+=1
            tempx =tempx + nets[0]
            nx+=1
        #print array(lattice)
        return lattice
    else:
        return False
# defining function for treatment of wrong input syntax. typ defines type of error, rest is standard print out.
def errorinput(typ):
    if typ==1:
        print "Too few arguments given to execute script!"
    if typ==2:
        print "Too many arguments given to execute script!"
    if typ==3:
        print "Unknown lattice type! Make sure lattice type is '1' or '2'!"
    print "\n Syntax: ligwrap.py <input><output>{rim}{netwitdth}{lattice type}"
    print "input: pdb or sph file containing ligand atoms or previous sphere set"
    print "output: should be name.sph and to this file the resulting spheres will be written"
    print "rim: width in A around original ligand which will contain the sphere lattice"
    print "netwitdth: approximate distance of the spheres"
    print "lattice type: 1 for simple cubic lattice (default), 2 for centered cubic lattice"
    print "if only two arguments are given rim, netwidth and lattice type will be set to default"
    print "if three arguments given, the netwidth and lattice type will be set to default value"
    print "if four arguments given, lattice type will be set to default"
          
	
#start of real script
#reading in coordinates and opening files
rim=1.5
lattc=2.0
latype=1
print str((sys.argv.__len__()-1))+" arguments given!"
if sys.argv.__len__()<3:
    errorinput(1)
    sys.exit()
elif sys.argv.__len__()==3:
    print "two arguments given, rim will be set to 1.5, netwidth will be set to 2.0"
    rim=1.5
    lattc=2.0
elif sys.argv.__len__()==4:
    print "three arguments given, netwidth will be set to 2.0"
    netwidth=2.0
    rim=float(sys.argv[3])
elif sys.argv.__len__()==5:
    print "four arguments given"
    rim=float(sys.argv[3])
    netwidth=float(sys.argv[4])
elif sys.argv.__len__()==6:
    latype=float(sys.argv[5])
    if latype==2.0:
        print "using centered cubic lattice type"
    elif latype==1.0:
        print "using simple cubic lattice type"
    else:
        errorinput(3)
        sys.exit()
    rim=float(sys.argv[3])
    netwidth=float(sys.argv[4])
else:
    errorinput(2)
    sys.exit()
data_file = open(sys.argv[1], 'r') #input file
out_file= open(sys.argv[2],'w') #output file containing .....(still under development)
print "Opened input and output files"

inpsplit=string.split(sys.argv[1],".") #split input file in prefix and suffix
dattype=inpsplit[1] 
primcord=[] #Varaible to store primary crystal lattice


atoms=readcoords(data_file)
[molcent,v,transc]=svdtrans(atoms)
box=boxcalc(transc,rim)
#print box
print type(latype)
primcord=array(latticegen(box,lattc,latype))


print "Calculated Sphere lattice"
#transform primary coordinates into array for convinience

#retransform and write out
retrans=dot(array(primcord),transpose(v)) #initialises variable and rotates coordinate axis back
molcent=array(molcent)
molcent=molcent.reshape([1,3]) #format to proper array
temp=dot(ones([primcord.shape[0],1]),molcent) #provide dummy for reshifting  
retrans=add(retrans,temp) #reshift into original position
n3=0
while n3<primcord.shape[0]:
	number="  "+str((n3+1))	
	out_file.write("  "+number[-3:]+" "+coordout(retrans[n3,0])+" "+coordout(retrans[n3,1])+" "+coordout(retrans[n3,2])+"   0.500    0 0  7\n")#write new coordintates embedded in the other information
	n3=n3+1
print "Retransformed lattice atoms into original coordinates and saved to out file" 
out_file.close()


#command="sph2pdb.py "+sys.argv[2]+" "+sys.argv[2]+".pdb"
sphconv(sys.argv[2])

sys.exit()
